

# torch.nn.Conv2d(
#     in_channels,   # 输入通道数（如RGB图像为3）
#     out_channels,  # 输出通道数（卷积核数量）
#     kernel_size,   # 卷积核大小（如3或(3, 3)）
#     stride=1,      # 步长（默认1）
#     padding=0,     # 填充（默认0，保持尺寸用 "same" 需计算）
#     dilation=1,    # 空洞卷积（默认1）
#     groups=1,      # 分组卷积（默认1）
#     bias=True      # 是否使用偏置项
# )
# 输入输出形状
# 输入：(batch_size, in_channels, height, width)
#
# 输出：(batch_size, out_channels, new_height, new_width)
#
# new_height = (height + 2*padding - kernel_size) // stride + 1
#
# new_width 同理。



# torch.nn.Linear(
#     in_features,   # 输入特征数
#     out_features,  # 输出特征数
#     bias=True      # 是否使用偏置项
# )


# (1) nn.CrossEntropyLoss（分类任务）
# output: 形状为 (batch_size, num_classes, ...)（未归一化的原始分数，即未经过Softmax）。
#
# 例如：(64, 10) 表示 batch_size=64，10 分类问题。
#
# target: 形状为 (batch_size, ...)，且是 类别索引（LongTensor，值范围 [0, num_classes-1]）。
#
# 例如：(64,) 或 (64, 1)，每个元素是 0 到 9 的整数。
#
# (2) nn.MSELoss（回归任务）
# output 和 target: 形状必须 完全相同（如 (batch_size, feature_dim)）。
#
# (3) nn.BCELoss（二分类）
# output: 形状 (batch_size, ...)，且经过 Sigmoid（值范围 [0, 1]）。
#
# target: 与 output 同形状，且值为 0 或 1。


# 可视化
# from torch.utils.tensorboard import SummaryWriter
# writer = SummaryWriter()
# writer.add_scalar('Loss/train', loss, epoch)

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn as nn
import torch





class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.fc1 = nn.Linear(32 * 26 * 26, 10)  # 根据输入尺寸计算

    def forward(self, x):
        temp = self.conv1(x)
        x = torch.relu(temp)
        x = x.view(x.size(0), -1)  # 展平
        x = self.fc1(x)
        return x

if __name__ == '__main__':

    # 数据增强与标准化
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485], std=[0.229])
    ])

    # 加载MNIST数据集
    train_data = datasets.MNIST(root='data/', train=True, transform=transform, download=True)
    train_loader = DataLoader(train_data, batch_size=64, shuffle=True)

    model = SimpleCNN()

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    for epoch in range(10):
        for batch_idx, (data, target) in enumerate(train_loader):
            print('batch:'+str(batch_idx))
            # data, target = data.to('cuda'), target.to('cuda')

            # 前向传播
            output = model(data)
            loss = criterion(output, target)

            # 反向传播与优化
            optimizer.zero_grad()  # 清除历史梯度！
            loss.backward()
            optimizer.step()
            print(model.conv1.weight.grad)  # 检查梯度是否更新
        print(f'Epoch {epoch}, Loss: {loss.item():.4f}')
    #
    # 验证模式
    # model.eval()
    # with torch.no_grad():  # 关闭梯度计算
    #     correct = 0
    #     for data, target in test_loader:
    #         output = model(data)
    #         pred = output.argmax(dim=1)
    #         correct += (pred == target).sum().item()
    #     print(f'Test Accuracy: {correct / len(test_data):.2f}')
    #
    # # 保存模型
    # torch.save(model.state_dict(), 'model.pth')
